"""
Session Manager - Central orchestrator for multi-website sessions.

This module provides the main interface for managing authentication sessions
across multiple websites, coordinating between authentication, storage, and
browser management components.
"""

import os
import json
import time
import logging
from typing import Dict, List, Optional, Set, Any
from datetime import datetime, timedelta
from pathlib import Path

from models import (
    WebsiteProfile, 
    UserSession, 
    GlobalConfiguration, 
    SessionStatus,
    create_user_session,
    get_default_global_config
)


class SessionManager:
    """
    Central orchestrator for multi-website session management.
    
    Coordinates authentication workflows, session lifecycle management,
    and synchronization across multiple websites and devices.
    """
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize the Session Manager.
        
        Args:
            config_path: Path to global configuration file
        """
        self.config_path = config_path or "global_config.json"
        self.config = self._load_config()
        self.logger = self._setup_logging()
        
        # Initialize storage directories
        self.sessions_dir = Path("sessions")
        self.profiles_dir = Path("website_profiles") 
        self.sessions_dir.mkdir(exist_ok=True)
        self.profiles_dir.mkdir(exist_ok=True)
        
        # In-memory caches
        self._website_profiles: Dict[str, WebsiteProfile] = {}
        self._active_sessions: Dict[str, UserSession] = {}
        
        # Load existing data
        self._load_website_profiles()
        self._load_active_sessions()
        
        self.logger.info("SessionManager initialized")
    
    def _load_config(self) -> GlobalConfiguration:
        """Load global configuration."""
        if os.path.exists(self.config_path):
            return GlobalConfiguration.load_from_file(self.config_path)
        else:
            # Create default config
            config = get_default_global_config()
            config.save_to_file(self.config_path)
            return config
    
    def _setup_logging(self) -> logging.Logger:
        """Setup logging configuration."""
        logger = logging.getLogger("SessionManager")
        logger.setLevel(getattr(logging, self.config.log_level))
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def _load_website_profiles(self):
        """Load website profiles from storage."""
        for profile_file in self.profiles_dir.glob("*.json"):
            try:
                with open(profile_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    profile = WebsiteProfile.from_dict(data)
                    self._website_profiles[profile.website_id] = profile
                    self.logger.debug(f"Loaded website profile: {profile.website_id}")
            except Exception as e:
                self.logger.error(f"Failed to load profile {profile_file}: {e}")
    
    def _load_active_sessions(self):
        """Load active sessions from storage."""
        for session_file in self.sessions_dir.glob("*.json"):
            try:
                with open(session_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    session = UserSession.from_dict(data)
                    
                    # Only load valid, non-expired sessions
                    if session.is_valid():
                        session_key = f"{session.username}_{session.website_id}"
                        self._active_sessions[session_key] = session
                        self.logger.debug(f"Loaded active session: {session_key}")
                    else:
                        self.logger.debug(f"Skipped invalid/expired session: {session.session_id}")
            except Exception as e:
                self.logger.error(f"Failed to load session {session_file}: {e}")
    
    def _save_website_profile(self, profile: WebsiteProfile):
        """Save website profile to storage."""
        profile_file = self.profiles_dir / f"{profile.website_id}.json"
        try:
            with open(profile_file, 'w', encoding='utf-8') as f:
                json.dump(profile.to_dict(), f, indent=2, ensure_ascii=False)
            self.logger.debug(f"Saved website profile: {profile.website_id}")
        except Exception as e:
            self.logger.error(f"Failed to save profile {profile.website_id}: {e}")
            raise
    
    def _save_session(self, session: UserSession):
        """Save session to storage."""
        session_file = self.sessions_dir / f"{session.session_id}.json"
        try:
            with open(session_file, 'w', encoding='utf-8') as f:
                json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
            self.logger.debug(f"Saved session: {session.session_id}")
        except Exception as e:
            self.logger.error(f"Failed to save session {session.session_id}: {e}")
            raise
    
    def add_website_profile(self, profile: WebsiteProfile) -> bool:
        """
        Add a new website profile.
        
        Args:
            profile: Website profile to add
            
        Returns:
            True if successful, False otherwise
        """
        try:
            if profile.website_id in self._website_profiles:
                self.logger.warning(f"Website profile already exists: {profile.website_id}")
                return False
            
            profile.update_timestamp()
            self._website_profiles[profile.website_id] = profile
            self._save_website_profile(profile)
            
            self.logger.info(f"Added website profile: {profile.website_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to add website profile: {e}")
            return False
    
    def update_website_profile(self, website_id: str, updates: Dict[str, Any]) -> bool:
        """
        Update an existing website profile.
        
        Args:
            website_id: Website identifier
            updates: Dictionary of fields to update
            
        Returns:
            True if successful, False otherwise
        """
        try:
            if website_id not in self._website_profiles:
                self.logger.error(f"Website profile not found: {website_id}")
                return False
            
            profile = self._website_profiles[website_id]
            
            # Update allowed fields
            for field, value in updates.items():
                if hasattr(profile, field):
                    setattr(profile, field, value)
            
            profile.update_timestamp()
            self._save_website_profile(profile)
            
            self.logger.info(f"Updated website profile: {website_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to update website profile: {e}")
            return False
    
    def remove_website_profile(self, website_id: str) -> bool:
        """
        Remove a website profile.
        
        Args:
            website_id: Website identifier
            
        Returns:
            True if successful, False otherwise
        """
        try:
            if website_id not in self._website_profiles:
                self.logger.warning(f"Website profile not found: {website_id}")
                return False
            
            # Remove associated sessions
            sessions_to_remove = [
                key for key in self._active_sessions.keys() 
                if key.endswith(f"_{website_id}")
            ]
            
            for session_key in sessions_to_remove:
                self.delete_session_by_key(session_key)
            
            # Remove profile
            del self._website_profiles[website_id]
            
            # Remove profile file
            profile_file = self.profiles_dir / f"{website_id}.json"
            if profile_file.exists():
                profile_file.unlink()
            
            self.logger.info(f"Removed website profile: {website_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to remove website profile: {e}")
            return False
    
    def get_website_profile(self, website_id: str) -> Optional[WebsiteProfile]:
        """Get website profile by ID."""
        return self._website_profiles.get(website_id)
    
    def list_website_profiles(self) -> List[WebsiteProfile]:
        """Get all website profiles."""
        return list(self._website_profiles.values())
    
    def create_session(self, website_id: str, username: str, auth_state: Dict[str, Any], 
                      extensions_data: Optional[bytes] = None) -> Optional[UserSession]:
        """
        Create a new user session.
        
        Args:
            website_id: Website identifier
            username: User account identifier
            auth_state: Authentication state data
            extensions_data: Browser extensions data
            
        Returns:
            Created session or None if failed
        """
        try:
            if website_id not in self._website_profiles:
                self.logger.error(f"Website profile not found: {website_id}")
                return None
            
            # Create session
            session = create_user_session(website_id, username, auth_state)
            session.extensions_data = extensions_data
            
            # Set expiry based on website configuration
            profile = self._website_profiles[website_id]
            expiry_hours = profile.session_config.timeout_hours
            session.expires_at = datetime.now() + timedelta(hours=expiry_hours)
            
            # Store session
            session_key = f"{username}_{website_id}"
            self._active_sessions[session_key] = session
            self._save_session(session)
            
            self.logger.info(f"Created session: {session_key}")
            return session
            
        except Exception as e:
            self.logger.error(f"Failed to create session: {e}")
            return None
    
    def get_session(self, username: str, website_id: str) -> Optional[UserSession]:
        """
        Get user session for specific website.
        
        Args:
            username: User account identifier
            website_id: Website identifier
            
        Returns:
            Session if found and valid, None otherwise
        """
        session_key = f"{username}_{website_id}"
        session = self._active_sessions.get(session_key)
        
        if session and session.is_valid():
            return session
        elif session and session.is_expired():
            # Remove expired session
            self.delete_session_by_key(session_key)
            
        return None
    
    def get_user_sessions(self, username: str) -> List[UserSession]:
        """
        Get all sessions for a specific user.
        
        Args:
            username: User account identifier
            
        Returns:
            List of active sessions for the user
        """
        user_sessions = []
        for session_key, session in self._active_sessions.items():
            if session.username == username and session.is_valid():
                user_sessions.append(session)
            elif session.is_expired():
                # Clean up expired session
                self.delete_session_by_key(session_key)
        
        return user_sessions
    
    def get_website_sessions(self, website_id: str) -> List[UserSession]:
        """
        Get all sessions for a specific website.
        
        Args:
            website_id: Website identifier
            
        Returns:
            List of active sessions for the website
        """
        website_sessions = []
        for session_key, session in self._active_sessions.items():
            if session.website_id == website_id and session.is_valid():
                website_sessions.append(session)
            elif session.is_expired():
                # Clean up expired session
                self.delete_session_by_key(session_key)
        
        return website_sessions
    
    def delete_session(self, username: str, website_id: str) -> bool:
        """
        Delete a specific session.
        
        Args:
            username: User account identifier
            website_id: Website identifier
            
        Returns:
            True if successful, False otherwise
        """
        session_key = f"{username}_{website_id}"
        return self.delete_session_by_key(session_key)
    
    def delete_session_by_key(self, session_key: str) -> bool:
        """Delete session by key."""
        try:
            if session_key in self._active_sessions:
                session = self._active_sessions[session_key]
                
                # Remove session file
                session_file = self.sessions_dir / f"{session.session_id}.json"
                if session_file.exists():
                    session_file.unlink()
                
                # Remove from memory
                del self._active_sessions[session_key]
                
                self.logger.info(f"Deleted session: {session_key}")
                return True
            else:
                self.logger.warning(f"Session not found: {session_key}")
                return False
                
        except Exception as e:
            self.logger.error(f"Failed to delete session: {e}")
            return False
    
    def clean_expired_sessions(self) -> int:
        """
        Remove all expired sessions.
        
        Returns:
            Number of sessions cleaned
        """
        expired_keys = []
        
        for session_key, session in self._active_sessions.items():
            if session.is_expired():
                expired_keys.append(session_key)
        
        cleaned_count = 0
        for key in expired_keys:
            if self.delete_session_by_key(key):
                cleaned_count += 1
        
        self.logger.info(f"Cleaned {cleaned_count} expired sessions")
        return cleaned_count
    
    def extend_session(self, username: str, website_id: str, hours: int = 24) -> bool:
        """
        Extend session expiry time.
        
        Args:
            username: User account identifier
            website_id: Website identifier
            hours: Hours to extend
            
        Returns:
            True if successful, False otherwise
        """
        try:
            session = self.get_session(username, website_id)
            if not session:
                self.logger.error(f"Session not found: {username}_{website_id}")
                return False
            
            session.extend_expiry(hours)
            self._save_session(session)
            
            self.logger.info(f"Extended session: {username}_{website_id} by {hours} hours")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to extend session: {e}")
            return False
    
    def get_session_summary(self) -> Dict[str, Any]:
        """
        Get summary of current session state.
        
        Returns:
            Dictionary with session statistics
        """
        total_sessions = len(self._active_sessions)
        total_websites = len(self._website_profiles)
        
        # Count sessions by website
        website_counts = {}
        user_counts = {}
        
        for session in self._active_sessions.values():
            website_counts[session.website_id] = website_counts.get(session.website_id, 0) + 1
            user_counts[session.username] = user_counts.get(session.username, 0) + 1
        
        return {
            "total_sessions": total_sessions,
            "total_websites": total_websites,
            "sessions_by_website": website_counts,
            "sessions_by_user": user_counts,
            "active_users": len(user_counts),
            "timestamp": datetime.now().isoformat()
        }
    
    def validate_sessions(self) -> Dict[str, List[str]]:
        """
        Validate all sessions and return status report.
        
        Returns:
            Dictionary with valid, expired, and invalid session lists
        """
        valid_sessions = []
        expired_sessions = []
        invalid_sessions = []
        
        for session_key, session in self._active_sessions.items():
            if session.is_expired():
                expired_sessions.append(session_key)
            elif session.is_valid():
                valid_sessions.append(session_key)
            else:
                invalid_sessions.append(session_key)
        
        return {
            "valid": valid_sessions,
            "expired": expired_sessions,
            "invalid": invalid_sessions
        }
    
    def get_config(self) -> GlobalConfiguration:
        """Get current global configuration."""
        return self.config
    
    def update_config(self, updates: Dict[str, Any]) -> bool:
        """
        Update global configuration.
        
        Args:
            updates: Dictionary of configuration updates
            
        Returns:
            True if successful, False otherwise
        """
        try:
            # Update configuration object
            config_dict = self.config.to_dict()
            config_dict.update(updates)
            
            self.config = GlobalConfiguration.from_dict(config_dict)
            self.config.save_to_file(self.config_path)
            
            self.logger.info("Updated global configuration")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to update configuration: {e}")
            return False